{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Self Correcting Query Engines - Evaluation & Retry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we showcase several advanced, self-correcting query engines.  \n",
    "They leverage the latest LLM's ability to evaluate its own output, and then self-correct to give better responses. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Uncomment to add your OpenAI API key\n",
    "# import os\n",
    "# os.environ['OPENAI_API_KEY'] = \"INSERT OPENAI KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Uncomment for debug level logging\n",
    "# import logging\n",
    "# import sys\n",
    "\n",
    "# logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)\n",
    "# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we ingest the document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.vector_store.base import VectorStoreIndex\n",
    "from llama_index.readers.file.base import SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\"../data/paul_graham/\").load_data()\n",
    "index = VectorStoreIndex.from_documents(documents)\n",
    "query = \"What did the author do growing up?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's what the response from the default query engine looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The author grew up writing essays, learning Italian, exploring Florence, painting people, learning about computers, attending RISD, living in a rent-stabilized apartment, building an online store builder, editing Lisp expressions, publishing essays online, writing essays, painting still life, working on spam filters, cooking for groups, and buying a building in Cambridge.\n"
     ]
    }
   ],
   "source": [
    "base_query_engine = index.as_query_engine()\n",
    "response = base_query_engine.query(query)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retry Query Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The retry query engine uses an evaluator to improve the response from a base query engine.  \n",
    "\n",
    "It does the following:\n",
    "1. first queries the base query engine, then\n",
    "2. use the evaluator to decided if the response passes.\n",
    "3. If the response passes, then return response,\n",
    "4. Otherwise, transform the original query with the evaluation result (query, response, and feedback) into a new query, \n",
    "5. Repeat up to max_retries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The author grew up writing essays, learning Italian, exploring Florence, painting people, working with computers, attending RISD, living in a rent-controlled apartment, building an online store builder, editing code, launching software, publishing essays online, writing essays, painting still life, working on spam filters, cooking for groups, and buying a building in Cambridge.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.query_engine import RetryQueryEngine\n",
    "from llama_index.evaluation import RelevancyEvaluator\n",
    "\n",
    "query_response_evaluator = RelevancyEvaluator()\n",
    "retry_query_engine = RetryQueryEngine(base_query_engine, query_response_evaluator)\n",
    "retry_response = retry_query_engine.query(query)\n",
    "print(retry_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Retry Source Query Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Source Retry modifies the query source nodes by filtering the existing source nodes for the query based on llm node evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The author grew up writing essays, learning Italian, exploring Florence, painting people, working with computers, attending RISD, living in a rent-stabilized apartment, building an online store builder, editing Lisp expressions, publishing essays online, writing essays, painting still life, working on spam filters, cooking for groups, and buying a building in Cambridge.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.query_engine import RetrySourceQueryEngine\n",
    "\n",
    "retry_source_query_engine = RetrySourceQueryEngine(\n",
    "    base_query_engine, query_response_evaluator\n",
    ")\n",
    "retry_source_response = retry_source_query_engine.query(query)\n",
    "print(retry_source_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Retry Guideline Query Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This module tries to use guidelines to direct the evaluator's behavior. You can customize your own guidelines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.evaluation.guideline_eval import GuidelineEvaluator, DEFAULT_GUIDELINES\n",
    "from llama_index.response.schema import Response\n",
    "from llama_index.indices.query.query_transform.feedback_transform import (\n",
    "    FeedbackQueryTransformation,\n",
    ")\n",
    "from llama_index.query_engine.retry_query_engine import (\n",
    "    RetryGuidelineQueryEngine,\n",
    ")\n",
    "\n",
    "# Guideline eval\n",
    "guideline_eval = GuidelineEvaluator(\n",
    "    guidelines=DEFAULT_GUIDELINES + \"\\nThe response should not be overly long.\\n\"\n",
    "    \"The response should try to summarize where possible.\\n\"\n",
    ")  # just for example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look like what happens under the hood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Guideline eval evaluation result: The response is too long and should be summarized. It should also include specific numbers or statistics when possible.\n",
      "Transformed query: Here is a previous bad answer.\n",
      "\n",
      "The author grew up writing essays, learning Italian, exploring Florence, painting people, learning about computers, attending RISD, living in a rent-stabilized apartment, building an online store builder, editing Lisp expressions, publishing essays online, writing essays, painting still life, working on spam filters, cooking for groups, and buying a building in Cambridge.\n",
      "Here is some feedback from the evaluator about the response given.\n",
      "The response is too long and should be summarized. It should also include specific numbers or statistics when possible.\n",
      "Now answer the question.\n",
      "\n",
      "What experiences did the author have growing up?\n"
     ]
    }
   ],
   "source": [
    "typed_response = response if isinstance(response, Response) else response.get_response()\n",
    "eval = guideline_eval.evaluate_response(query, typed_response)\n",
    "print(f\"Guideline eval evaluation result: {eval.feedback}\")\n",
    "\n",
    "feedback_query_transform = FeedbackQueryTransformation(resynthesize_query=True)\n",
    "transformed_query = feedback_query_transform.run(query, {\"evaluation\": eval})\n",
    "print(f\"Transformed query: {transformed_query.query_str}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now let's run the full query engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The author gained a wide range of skills and experiences growing up, from creative pursuits such as painting and writing, to technical skills such as coding and launching software. They also gained an understanding of Italian language and culture, explored the city of Florence, and gained experience living in a rent-stabilized apartment. These experiences enabled the author to develop the skills necessary to launch a successful online store builder and publish essays online, which allowed them to reach a wider audience and gain recognition for their work. Specifically, they gained an understanding of the web infrastructure, the ability to write code in Lisp, and the ability to market their products and services. They also developed an understanding of the importance of working on projects that weren't prestigious, as well as the ability to cook for large groups of people. Finally, they gained the confidence to take risks and pursue their passions, which ultimately led to their success.\n"
     ]
    }
   ],
   "source": [
    "retry_guideline_query_engine = RetryGuidelineQueryEngine(\n",
    "    base_query_engine, guideline_eval, resynthesize_query=True\n",
    ")\n",
    "retry_guideline_response = retry_guideline_query_engine.query(query)\n",
    "print(retry_guideline_response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
